import math
import pandas as pd
import torch
from torch import nn
from d2l import torch as d2l


# @save
class PositionWiseFFN(nn.Module):
    """基于位置的前馈网络"""

    def __init__(self, ffn_num_input, ffn_num_hiddens, ffn_num_outputs,
                 **kwargs):
        super(PositionWiseFFN, self).__init__(**kwargs)
        self.dense1 = nn.Linear(ffn_num_input, ffn_num_hiddens)
        self.relu = nn.ReLU()
        self.dense2 = nn.Linear(ffn_num_hiddens, ffn_num_outputs)

    def forward(self, X):
        return self.dense2(self.relu(self.dense1(X)))


# 输入2，3，4 输出 2，3，8 最后一维度和ffn输入第一个参数相同，输出将最后一个维度改变为ffn的第三个参数大小。


#
# # 组间对应元素标准化
# X1 = torch.tensor([[1, 300], [2, 200], [3, 100]], dtype=torch.float32)
# bn = nn.BatchNorm1d(2)
# print(bn(X1))
# # 组内所有元素标准化
# X2 = torch.tensor([[1, 2, 3], [300, 200, 100]], dtype=torch.float32)
# ln = nn.LayerNorm(3)
# print(ln(X2))
# # 即使提高到2维，也是在2维组内进行运算。
# X3 = torch.tensor([[[1, 2], [3, 4], [5, 6]], [[100, 300], [200, 400], [500, 600]]], dtype=torch.float32)
# ln = nn.LayerNorm([3, 2])
# print(ln(X3))


# @save
class AddNorm(nn.Module):
    """残差连接后进行层规范化"""

    def __init__(self, normalized_shape, dropout, **kwargs):
        super(AddNorm, self).__init__(**kwargs)
        self.dropout = nn.Dropout(dropout)
        self.ln = nn.LayerNorm(normalized_shape)

    def forward(self, X, Y):
        return self.ln(self.dropout(Y) + X)


# 在24，100组内进行规范化操作，样本之间不做规范化
add_norm = AddNorm([100, 24], 0.5)
add_norm.eval()

# 逐位前馈网络
ffn = PositionWiseFFN(24, 64, 24)
ffn.eval()

# 输入的input有可能是原始输入 或者多头注意力输出 前馈网络有可能是
input = torch.ones((32, 100, 24))
print("单层编码器编码前信息", input.shape)

print("逐位前馈网路处理融合后的信息", ffn(input).shape)

# posit_out 是对输入经过逐位前馈网络的输出
posit_out = ffn(input)
# 逐位前馈网络的输出和刚才的输入进行残差连接规范化过程
output = add_norm(input, posit_out)
print("残差连接后进行层规范化，用于将原始特征和经过多头注意力处理后的特征进行融合", output.shape)


# @save
class EncoderBlock(nn.Module):
    """Transformer编码器块"""

    def __init__(self, key_size, query_size, value_size, num_hiddens,
                 norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,
                 dropout, use_bias=False, **kwargs):
        super(EncoderBlock, self).__init__(**kwargs)
        # 区块包含一个多头注意力头，一个残差连接标准化层，一个前馈神经网络，和第二个残差连接标准化层
        self.attention = d2l.MultiHeadAttention(
            key_size, query_size, value_size, num_hiddens, num_heads, dropout,
            use_bias)
        self.addnorm1 = AddNorm(norm_shape, dropout)
        self.ffn = PositionWiseFFN(
            ffn_num_input, ffn_num_hiddens, num_hiddens)
        self.addnorm2 = AddNorm(norm_shape, dropout)

    def forward(self, X, valid_lens):
        # 多头自注意力
        Y = self.addnorm1(X, self.attention(X, X, X, valid_lens))
        return self.addnorm2(Y, self.ffn(Y))


x_valid_lens = torch.tensor([100] * 32)
"""
key_size 24
query_size 24
value_size 24
num_hiddens 24        
norm_shape [100, 24] 
ffn_num_input 24
ffn_num_hiddens 48
num_heads 8
dropout 0.5
use_bias=False
"""
key_size = 24
query_size = 24
value_size = 24
num_hiddens = 24
norm_shape = [100, 24]
ffn_num_input = 24
ffn_num_hiddens = 48
num_heads = 8
dropout = 0.5

encoder_blk = EncoderBlock(
    key_size, query_size, value_size, num_hiddens, norm_shape,
    ffn_num_input, ffn_num_hiddens,
    num_heads, dropout)
encoder_blk.eval()
print("单层编码器编码后信息", encoder_blk(input, x_valid_lens).shape)


class TransformerEncoder(d2l.Encoder):
    """Transformer编码器"""

    def __init__(self, vocab_size, key_size, query_size, value_size,
                 num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens,
                 num_heads, num_layers, dropout, use_bias=False, **kwargs):
        super(TransformerEncoder, self).__init__(**kwargs)
        self.num_hiddens = num_hiddens
        self.embedding = nn.Embedding(vocab_size, num_hiddens)
        self.pos_encoding = d2l.PositionalEncoding(num_hiddens, dropout)
        self.blks = nn.Sequential()
        for i in range(num_layers):
            self.blks.add_module("block" + str(i),
                                 EncoderBlock(key_size, query_size, value_size, num_hiddens,
                                              norm_shape, ffn_num_input, ffn_num_hiddens,
                                              num_heads, dropout, use_bias))

    def forward(self, X, valid_lens, *args):
        # 因为位置编码值在-1和1之间，
        # 因此嵌入值乘以嵌入维度的平方根进行缩放，
        # 然后再与位置编码相加。
        # print("X 在进行词嵌入且位置编码之前",X.shape)
        X = self.pos_encoding(self.embedding(X) * math.sqrt(self.num_hiddens))
        # print("X 在进行词嵌入并进行位置编码之后", X.shape)
        # 初始化两个变量记录注意力权重
        self.attention_weights = [None] * len(self.blks)
        # print("self.attention_weights",self.attention_weights)
        for i, blk in enumerate(self.blks):
            # 对 X进行区块（EncoderBlock）运算
            X = blk(X, valid_lens)
            # 循环更新注意力权重 两个大小为 [64, 12, 12]的注意力权重
            # print("注意力权重",blk.attention.attention.attention_weights.shape)
            self.attention_weights[
                i] = blk.attention.attention.attention_weights
        return X


vocab_size = 20000
num_layers = 6
# 将32批次每一个批次120个词传入到模型。
# 输出为
# batch_size, vocab_lens
X = torch.ones((32, 100), dtype=torch.long)
encoder = TransformerEncoder(
    # vocab_size, key_size, query_size, value_size, num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens, num_heads, num_layers, dropout)
    vocab_size, key_size, query_size, value_size, num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,
    num_layers, dropout)
encoder.eval()
# torch.Size([32, 100, 24])
print("针对原始数据的编码输出", encoder(X, x_valid_lens).shape)

# d2l.DotProductAttention
# torch.Size([batch_size * heads , vocab_lens, vocab_lens])
# torch.Size([256, 100, 100])
# print("第一层信息", encoder.attention_weights[0].shape)  torch.Size([256, 100, 100])
# print("第二层信息", encoder.attention_weights[1].shape)  torch.Size([256, 100, 100])
class DecoderBlock(nn.Module):
    """解码器中第i个块"""

    def __init__(self, key_size, query_size, value_size, num_hiddens,
                 norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,
                 dropout, i, **kwargs):
        super(DecoderBlock, self).__init__(**kwargs)
        self.i = i
        self.attention1 = d2l.MultiHeadAttention(
            key_size, query_size, value_size, num_hiddens, num_heads, dropout)
        self.addnorm1 = AddNorm(norm_shape, dropout)
        self.attention2 = d2l.MultiHeadAttention(
            key_size, query_size, value_size, num_hiddens, num_heads, dropout)
        self.addnorm2 = AddNorm(norm_shape, dropout)
        self.ffn = PositionWiseFFN(ffn_num_input, ffn_num_hiddens,
                                   num_hiddens)
        self.addnorm3 = AddNorm(norm_shape, dropout)

    def forward(self, X, state):
        # 编码器输出，编码器输出的有效长度
        enc_outputs, enc_valid_lens = state[0], state[1]
        batch_size, num_steps, _ = X.shape
        # dec_valid_lens的开头:(batch_size,num_steps),
        # 其中每一行是[1,2,...,num_steps]
        dec_valid_lens = torch.arange(
            1, num_steps + 1, device=X.device).repeat(batch_size, 1)

        # 自注意力,后一个特征是前面几个特征的加权求和，是针对自身的特征修正
        X2 = self.attention1(X, X, X, dec_valid_lens)
        Y = self.addnorm1(X, X2)
        # 编码器－解码器注意力。
        # enc_outputs的开头:(batch_size,num_steps,num_hiddens)
        Y2 = self.attention2(Y, enc_outputs, enc_outputs, enc_valid_lens)
        Z = self.addnorm2(Y, Y2)
        return self.addnorm3(Z, self.ffn(Z)), state


#
# class DecoderBlock(nn.Module):
#     """解码器中第i个块"""
#
#     def __init__(self, key_size, query_size, value_size, num_hiddens,
#                  norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,
#                  dropout, i, **kwargs):
#         super(DecoderBlock, self).__init__(**kwargs)
#         self.i = i
#         self.attention1 = d2l.MultiHeadAttention(
#             key_size, query_size, value_size, num_hiddens, num_heads, dropout)
#         self.addnorm1 = AddNorm(norm_shape, dropout)
#         self.attention2 = d2l.MultiHeadAttention(
#             key_size, query_size, value_size, num_hiddens, num_heads, dropout)
#         self.addnorm2 = AddNorm(norm_shape, dropout)
#         self.ffn = PositionWiseFFN(ffn_num_input, ffn_num_hiddens,
#                                    num_hiddens)
#         self.addnorm3 = AddNorm(norm_shape, dropout)
#
#     def forward(self, X, state):
#         # 编码器输出，编码器输出的有效长度
#         enc_outputs, enc_valid_lens = state[0], state[1]
#         # 训练阶段，输出序列的所有词元都在同一时间处理，
#         # 因此state[2][self.i]初始化为None。
#         # 预测阶段，输出序列是通过词元一个接着一个解码的，
#         # 因此state[2][self.i]包含着直到当前时间步第i个块解码的输出表示
#         if state[2][self.i] is None:
#             key_values = X
#         else:
#             key_values = torch.cat((state[2][self.i], X), axis=1)
#         state[2][self.i] = key_values
#         if self.training:
#             # 32        120       128
#             batch_size, num_steps, _ = X.shape
#             # dec_valid_lens的开头:(batch_size,num_steps),
#             # 其中每一行是[1,2,...,num_steps]
#             dec_valid_lens = torch.arange(
#                 1, num_steps + 1, device=X.device).repeat(batch_size, 1)
#             # batch_size,num_steps
#             # torch.Size([32, 120]) dec_valid_lens
#         else:
#             dec_valid_lens = None
#         # print("判断X, key_values是否相等。。。。。。。。。。。。。。。。",torch.equal(X, key_values))
#         # 自注意力
#         X2 = self.attention1(X, key_values, key_values, dec_valid_lens)
#         Y = self.addnorm1(X, X2)
#         # print("判断Y, enc_outputs是否相等。。。。。。。。。。。。。。。。", torch.equal(Y, enc_outputs))
#         Y2 = self.attention2(Y, enc_outputs, enc_outputs, enc_valid_lens)
#         Z = self.addnorm2(Y, Y2)
#
#         return self.addnorm3(Z, self.ffn(Z)), state


decoder_blk0 = DecoderBlock(24, 24, 24, 24, [150, 24], 24, 48, 8, 0.5, 0)
decoder_blk0.eval()
decoder_blk1 = DecoderBlock(24, 24, 24, 24, [150, 24], 24, 48, 8, 0.5, 1)
decoder_blk1.eval()
X = torch.ones((32, 100), dtype=torch.long)
state = [encoder(X, x_valid_lens), x_valid_lens, [None, None]]

print("32批次100个token的新词向量表示", state[0].shape)
Y = torch.ones((32, 150, 24))
# Y 是32批次 长度为 60 的词向量表示
# state包含了3个元素
# 1. 编码器的输出表示
# 2. 编码器的有效长度
# 3. 解码器中第i个块的输出表示
X0, state = decoder_blk0(Y, state)


# class TransformerDecoder(d2l.AttentionDecoder):
#     def __init__(self, vocab_size, key_size, query_size, value_size,
#                  num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens,
#                  num_heads, num_layers, dropout, **kwargs):
#         super(TransformerDecoder, self).__init__(**kwargs)
#         self.num_hiddens = num_hiddens
#         self.num_layers = num_layers
#         self.embedding = nn.Embedding(vocab_size, num_hiddens)
#         self.pos_encoding = d2l.PositionalEncoding(num_hiddens, dropout)
#         self.blks = nn.Sequential()
#         for i in range(num_layers):
#             self.blks.add_module("block" + str(i),
#                                  DecoderBlock(key_size, query_size, value_size, num_hiddens,
#                                               norm_shape, ffn_num_input, ffn_num_hiddens,
#                                               num_heads, dropout, i))
#         self.dense = nn.Linear(num_hiddens, vocab_size)
#
#     def init_state(self, enc_outputs, enc_valid_lens, *args):
#         # print("初始化状态",enc_valid_lens)
#         # torch.Size([16, 12, 32]) tensor([4, 5, 5, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4])
#         return [enc_outputs, enc_valid_lens, [None] * self.num_layers]
#
#     def forward(self, X, state):
#         X = self.pos_encoding(self.embedding(X) * math.sqrt(self.num_hiddens))
#         self._attention_weights = [[None] * len(self.blks) for _ in range(2)]
#         for i, blk in enumerate(self.blks):
#             X, state = blk(X, state)
#             # 解码器自注意力权重
#             self._attention_weights[0][
#                 i] = blk.attention1.attention.attention_weights
#             # “编码器－解码器”自注意力权重
#             self._attention_weights[1][
#                 i] = blk.attention2.attention.attention_weights
#         return self.dense(X), state
#
#     @property
#     def attention_weights(self):
#         return self._attention_weights


class TransformerDecoder(d2l.AttentionDecoder):
    def __init__(self, vocab_size, key_size, query_size, value_size,
                 num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens,
                 num_heads, num_layers, dropout, **kwargs):
        super(TransformerDecoder, self).__init__(**kwargs)
        self.num_hiddens = num_hiddens
        self.num_layers = num_layers
        self.embedding = nn.Embedding(vocab_size, num_hiddens)
        self.pos_encoding = d2l.PositionalEncoding(num_hiddens, dropout)
        self.blks = nn.Sequential()
        for i in range(num_layers):
            self.blks.add_module("block" + str(i),
                                 DecoderBlock(key_size, query_size, value_size, num_hiddens,
                                              norm_shape, ffn_num_input, ffn_num_hiddens,
                                              num_heads, dropout, i))
        self.dense = nn.Linear(num_hiddens, vocab_size)

    def init_state(self, enc_outputs, enc_valid_lens, *args):
        # enc_outputs记录了编码器输出结果  enc_valid_lens记录了编码器有效长度
        self.seqX = None
        return [enc_outputs, enc_valid_lens]

    def forward(self, X, state):
        # 这里面的X指的是解码器输入，在这里特指加了<bos>标识的解码序列
        # 如果不是处于训练阶段
        if not self.training:
            self.seqX = X if self.seqX is None else torch.cat((self.seqX, X), dim=1)
            X = self.seqX

        X = self.pos_encoding(self.embedding(X) * math.sqrt(self.num_hiddens))
        self._attention_weights = [[None] * len(self.blks) for _ in range(2)]
        for i, blk in enumerate(self.blks):
            X, state = blk(X, state)
            # 解码器自注意力权重
            self._attention_weights[0][
                i] = blk.attention1.attention.attention_weights
            # “编码器－解码器”自注意力权重
            self._attention_weights[1][
                i] = blk.attention2.attention.attention_weights

        if not self.training:
            # print("预测过程",self.dense(X)[:, -1:, :].shape)
            return self.dense(X)[:, -1:, :], state
        # print("训练过程", self.dense(X).shape)
        return self.dense(X), state

    @property
    def attention_weights(self):
        return self._attention_weights


# 10.7.6. 训练
num_hiddens, num_layers, dropout, batch_size, num_steps = 32, 2, 0.1, 64, 10
lr, num_epochs, device = 0.005, 3, d2l.try_gpu()
ffn_num_input, ffn_num_hiddens, num_heads = 32, 64, 4
key_size, query_size, value_size = 32, 32, 32
norm_shape = [32]

train_iter, src_vocab, tgt_vocab = d2l.load_data_nmt(batch_size, num_steps)

encoder = TransformerEncoder(
    len(src_vocab), key_size, query_size, value_size, num_hiddens,
    norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,
    num_layers, dropout)
decoder = TransformerDecoder(
    len(tgt_vocab), key_size, query_size, value_size, num_hiddens,
    norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,
    num_layers, dropout)
net = d2l.EncoderDecoder(encoder, decoder)

# ##########################################
# # 模式一，seq2seq模式
# X = torch.ones((32, 100), dtype=torch.long)
# Y = torch.ones((32, 100), dtype=torch.long)
# x_valid_lens = torch.tensor(32 * [100])
# # Y_valid_len = torch.tensor(1 * [4])
# bos = torch.tensor([tgt_vocab['<bos>']] * Y.shape[0],
#                    device=device).reshape(-1, 1)
# # 将翻译后的数据进行错位增加一个开头标识
# dec_input = d2l.concat([bos, Y[:, :-1]], 1)  # Teacher forcing
# # 将源句子，带有强制教学信息的目的句子和X有效长度信息输入到模型中
# print("开始进入transformers ########################")
# print("源句子", X.shape)
# print("目的句子", Y.shape)
# print("前面加上bos前缀且舍弃最后一个词的 Y ", dec_input.shape)
# Y_hat, _ = net(X, dec_input, x_valid_lens)
# print(Y_hat.shape)  # 都是 [32,100]
# print("结束 transformers操作 ########################")
# ##########################################

d2l.train_seq2seq(net, train_iter, lr, num_epochs, tgt_vocab, device)
engs = ['go .', "i lost .", 'he\'s calm .', 'i\'m home .','i want to eat you .']
fras = ['va !', 'j\'ai perdu .', 'il est calme .', 'je suis chez moi .','je veux te manger .']
for eng, fra in zip(engs, fras):
    translation, dec_attention_weight_seq = d2l.predict_seq2seq(
        net, eng, src_vocab, tgt_vocab, num_steps, device, True)
    print(f'{eng} => {translation}, ',
          f'bleu {d2l.bleu(translation, fra, k=2):.3f}')

# enc_attention_weights = torch.cat(net.encoder.attention_weights, 0).reshape((num_layers, num_heads,
#                                                                              -1, num_steps))
# print(enc_attention_weights.shape)
